# %load AI/code/dataloader/data_dataloader.py
# 制作数据的数据导入器
'''
数据输入：numpy
类型：
默认值：

数据输出：Tensor
类型：
默认值：
'''
train_data_transform = transforms.Compose([
    transforms.Resize((84, 84)),
    transforms.CenterCrop(84),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
val_data_transform = transforms.Compose([
    transforms.Resize((84, 84)),
    transforms.CenterCrop(84),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
test_data_transform = transforms.Compose([
    transforms.Resize((84, 84)),
    transforms.CenterCrop(84),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 设置batch大小
train_batch_size = 256
val_batch_size = 256
test_batch_size = 256

# DataSet类 - （目录中图片导入）
train_dataset = datasets.ImageFolder(root=train_data_path, transform=train_data_transform)
val_dataset = datasets.ImageFolder(root=val_data_path, transform=val_data_transform)
val_dataset = datasets.ImageFolder(root=test_data_path, transform=val_data_transform)

# 数据加载器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True, num_workers=0)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=val_batch_size, shuffle=True, num_workers=0)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False, num_workers=0)
